# -*- coding: utf-8 -*-
#  带check 处理
# @Time    ： 2021/12/1  23:08
# @Auther  ： 徐良
# File     :  baseTreeWidget.py
from PySide6.QtWidgets import QTreeWidget
from PySide6.QtCore import Qt


class BaseTreeWidget(QTreeWidget):

    def __init__(self, parent=None):
        super(BaseTreeWidget, self).__init__(parent)
        self.setColumnCount(2)
        self.itemChanged.connect(self.treeItemChange)

    def treeItemChange(self, item, column):
        """
        item check发生变化处理方法
        :param item:
        :param column:
        :return:
        """
        self.itemChanged.disconnect(self.treeItemChange)  # 断开槽

        self.doTreeItemChange(item, column)  # 处理子节点 如果父节点选中，则子节点全部选中
        self.updateParentItem(item, column)  # 处理父节点

        self.itemChanged.connect(self.treeItemChange)  # 连接槽

    def treeExpandAll(self):
        self.expandAll()

    def treeCollapse(self):
        self.collapseAll()

    def doTreeItemChange(self, item, column):
        # 选中时 如果父节点选中，则子节点全部选中
        if (Qt.Checked == item.checkState(0)):
            count = item.childCount()
            if count > 0:
                for i in range(0, count):
                    # 子节点页被选中
                    item.child(i).setCheckState(0, Qt.Checked)
                    if item.childCount() > 0:
                        self.doTreeItemChange(item.child(i), column)
        else:
            # 未选中 如果父节点未选中，则子节点全部未选中
            count = item.childCount()
            if count > 0:
                for i in range(0, count):
                    # 子节点页被选中
                    item.child(i).setCheckState(0, Qt.Unchecked)
                    if item.childCount() > 0:
                        self.doTreeItemChange(item.child(i), column)

    def updateParentItem(self, item, column):
        parentitem = item.parent()
        if parentitem == None:
            return

        nChildCount = parentitem.childCount()
        nSelected = 0  # 选择状态
        nPartSelected = 0  # 部分选择状态

        for i in range(0, nChildCount):
            if parentitem.child(i).checkState(0) == Qt.Checked:
                nSelected += 1
            elif parentitem.child(i).checkState(0) == Qt.PartiallyChecked:
                nPartSelected += 1

        # 如有节点被选中，但选中的节点个数小于子节点数，或者子节点是半选中状态，则改节点为半选中
        if nSelected > 0 and nSelected < nChildCount or nPartSelected > 0:
            parentitem.setCheckState(0, Qt.PartiallyChecked)
        elif nSelected == nChildCount:
            parentitem.setCheckState(0, Qt.Checked)
        else:
            parentitem.setCheckState(0, Qt.Unchecked)

        self.updateParentItem(parentitem, column)

    def setNodeCheckAll(self):
        """
        设置全部节点选中
        :return:
        """

        def setCheckAll(node):
            count = node.childCount()
            for j in range(0, count):
                node.child(j).setCheckState(0, Qt.Checked)
                if node.child(j).childCount() > 0:
                    setCheckAll(node.child(j))

        n = self.topLevelItemCount()  # 获取根节点数量
        for i in range(0, n):
            item = self.topLevelItem(i)
            item.setCheckState(0, Qt.Checked)
            setCheckAll(item)

    def setNodeUnCheckAll(self):
        """
        全部节点设置未选中
        :return:
        """
        n = self.topLevelItemCount()  # 获取根节点数量
        for i in range(0, n):
            item = self.topLevelItem(i)
            item.setCheckState(0, Qt.Unchecked)

    def getAllNode(self):
        """
        获取所有节点
        :return: list[node:QtreeWidgetItem]
        """

        def getChild(item, itemList):
            count = item.childCount()
            for i in range(0, count):
                subItem = item.child(i)
                itemList.append(subItem)
                if subItem.childCount():
                    getChild(subItem, itemList)

        itemList = []
        n = self.topLevelItemCount()  # 获取根节点数量
        for i in range(0, n):
            item = self.topLevelItem(i)
            itemList.append(item)
            getChild(item, itemList)

        return itemList

    def getNode(self):
        """
        获取当前节点及其所有的子节点
        :return:
        """
        pass

    def getAllCheckedNode(self):
        '''
        获取所有已经选中的节点
        :return: list[node:QtreeWidgetItem]
        '''
        '''
        获取所有选中的叶子节点
        :return: list[node:QtreeWidgetItem]
        '''

        def getChild(node, itemList):
            icount = node.childCount()
            for i in range(0, icount):
                subitem = node.child(i)
                if subitem.checkState(0) == Qt.Checked or subitem.checkState(0) == Qt.PartiallyChecked:
                    itemList.append(subitem)
                    if subitem.childCount() > 0:
                        getChild(subitem, itemList)

        itemList = []
        n = self.topLevelItemCount()  # 获取根节点数量
        for i in range(0, n):
            item = self.topLevelItem(i)
            if item.checkState(0) == Qt.Checked or item.checkState(0) == Qt.PartiallyChecked:
                itemList.append(item)  # 是叶子节点
                if item.childCount() > 0:
                    getChild(item, itemList)

        return itemList

    def getAllCheckedLeaf(self):
        """
        获取所有选中的叶子节点
        :return: list[node:QtreeWidgetItem]
        """

        def getChild(node, itemList):
            icount = node.childCount()
            for i in range(0, icount):
                subitem = node.child(i)
                if subitem.checkState(0) == Qt.Checked or subitem.checkState(0) == Qt.PartiallyChecked:
                    if subitem.childCount() > 0:
                        getChild(subitem, itemList)
                    else:
                        itemList.append(subitem)

        itemList = []
        n = self.topLevelItemCount()  # 获取根节点数量
        for i in range(0, n):
            item = self.topLevelItem(i)
            if item.checkState(0) == Qt.Checked or item.checkState(0) == Qt.PartiallyChecked:
                if item.childCount() > 0:
                    getChild(item, itemList)
                else:
                    itemList.append(item)  # 是叶子节点

        return itemList
